﻿/*
* ////////////////////////////////////////////////
*
* 开发者：AZPEP
* 友情链接：
*	-	B站：		零分的考卷			https://space.bilibili.com/257391258
*	-	Github：	AZPEP-BiliBili		https://github.com/AZPEP-BiliBili
*	-	Gitee：		AZPEP_BiliBili		https://gitee.com/AZPEP_BiliBili
*
* CopyRight AZPEP
*
* ////////////////////////////////////////////////
*/

#include "pch.h"
#include "framework.h"

#include "GetEsp.h"
#include "GetGPT.h"

#include "includes/MountESP.h"

/// <summary>
///  获取磁盘上所有空闲的盘符
/// </summary>
/// <returns>成功GetLastError()=0，失败GetLastError()为非零值</returns>
std::vector<WCHAR> GetAvaiableLetter();

/// <summary>
/// 挂载计算机上全部的EFI分区
/// </summary>
/// <param name="LetterList">指向存放挂载分区盘符的列表的指针，可为nullptr</param>
/// <param name="ReMountPartition">是否重新挂载，默认FALSE</param>
/// <returns>成功返回0，LetterList指向存放挂载分区盘符的列表的指针；没有足够的盘符返回ERROR_NO_MORE_AVAIABLE_LETTER，失败返回GetLastError()得到的错误代码；若返回代码为非零值，则LetterList为空</returns>
DWORD AstMountEsp_MountEfiPartitionW(
	_Out_ std::vector<WCHAR>* LetterList,
	_In_ BOOL ReMountPartition
) {
	DWORD dStatus;
	//指定ReMountPartition就卸载全部EFI分区
	if (ReMountPartition) {
		dStatus = AstMountEsp_UnMountEfiPartitionW(nullptr);
		if (dStatus != 0) return dStatus;
	}

	//获取全部EFI分区
	std::vector<PartitionInformationW*> Esp_PIL;
	dStatus = MountEsp_GetEfiPartitionInfoW(&Esp_PIL);
	if (dStatus != 0) return dStatus;

	//获取可用盘符
	std::vector<WCHAR> AvaiableLetter = GetAvaiableLetter();

	//没有更多盘符，返回ERROR_NO_MORE_AVAIABLE_LETTER
	if (AvaiableLetter.size() < Esp_PIL.size()) { return ERROR_NO_MORE_AVAIABLE_LETTER; }

	//挂载EFI分区
	for (int i = 0; i < Esp_PIL.size(); i++) {
		BOOL bStatus = AstMountEsp_MountDriveW(
			AvaiableLetter[i],
			Esp_PIL[i]->PartitionPath
		);

		if (!bStatus) { LetterList->clear(); return GetLastError(); }
		if (LetterList != nullptr)LetterList->push_back(AvaiableLetter[i]);
	}

	return 0;
}

/// <summary>
/// 挂载计算机指定磁盘上全部EFI分区
/// </summary>
/// <param name="DiskNumber">将要执行EFI分区挂载操作的磁盘号</param>
/// <param name="LetterList">指向存放挂载分区盘符的列表的指针，可为nullptr</param>
/// <param name="ReMountPartition">是否重新挂载，默认FALSE</param>
/// <returns>成功返回0，LetterList指向存放挂载分区盘符的列表的指针，磁盘不存在返回ERROR_DRIVE_NOT_EXIST，没有足够的盘符返回ERROR_NO_MORE_AVAIABLE_LETTER，失败返回GetLastError()得到的错误代码；若返回代码为非零值，则LetterList为空</returns>
DWORD AstMountEsp_MountEfiPartitionW(
	_In_ DWORD DiskNumber,
	_Out_ std::vector<WCHAR>* LetterList,
	_In_ BOOL ReMountPartition
) {
	DWORD dStatus;

	//尝试打开磁盘
	std::wstringstream wss_PhysicalDrivePath;
	wss_PhysicalDrivePath << L"\\\\.\\PhysicalDrive";
	wss_PhysicalDrivePath << DiskNumber;

	HANDLE hDevice = INVALID_HANDLE_VALUE; //硬盘句柄
	//打开磁盘
	hDevice = CreateFile(
		wss_PhysicalDrivePath.str().c_str(),
		GENERIC_READ,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL
	);
	if (hDevice == INVALID_HANDLE_VALUE) return ERROR_DRIVE_NOT_EXIST;

	//指定ReMountPartition就卸载全部EFI分区
	if (ReMountPartition) {
		dStatus = AstMountEsp_UnMountEfiPartitionW(DiskNumber, nullptr);
		if (dStatus != 0) return dStatus;
	}

	//获取全部EFI分区
	std::vector<PartitionInformationW*> Esp_PIL;
	dStatus = MountEsp_GetEfiPartitionInfoW(&Esp_PIL);
	if (dStatus != 0) return dStatus;

	//获取可用盘符
	std::vector<WCHAR> AvaiableLetter = GetAvaiableLetter();

	//没有更多盘符，返回ERROR_NO_MORE_AVAIABLE_LETTER
	if (AvaiableLetter.size() < Esp_PIL.size()) { return ERROR_NO_MORE_AVAIABLE_LETTER; }

	//挂载EFI分区
	for (int i = 0; i < Esp_PIL.size(); i++) {
		if (Esp_PIL[i]->DiskNumber != DiskNumber) continue;

		BOOL bStatus = AstMountEsp_MountDriveW(
			AvaiableLetter[i],
			Esp_PIL[i]->PartitionPath
		);

		if (!bStatus) { LetterList->clear(); return GetLastError(); }
		if (LetterList != nullptr)LetterList->push_back(AvaiableLetter[i]);
	}

	return 0;
}

/// <summary>
/// 卸载计算机上全部EFI分区，已经卸载的不会重复卸载
/// </summary>
/// <param name="LetterList">指向卸载的EFI分区盘符列表的指针，可为nullptr</param>
/// <returns>成功返回0，LetterList存储卸载的EFI分区盘符，失败返回GetLastError()返回的错误代码</returns>
DWORD AstMountEsp_UnMountEfiPartitionW(
	_Out_ std::vector<WCHAR>* LetterList
) {
	DWORD dStatus;

	//获取全部EFI分区信息
	std::vector<PartitionInformationW*> vec_Efi_PIW;
	dStatus = MountEsp_GetEfiPartitionInfoW(&vec_Efi_PIW);
	if (dStatus != 0) { return dStatus; }	//获取失败，返回错误代码

	//遍历列表
	for (auto iter = vec_Efi_PIW.begin(); iter < vec_Efi_PIW.end(); iter++) {
		//查看分区是否已挂载
		if ((*iter)->IsMounted) {
			//将卸载分区压入LetterList
			if (LetterList != nullptr) LetterList->push_back((*iter)->Letter);

			//如果挂载，就卸载分区
			BOOL bStatus = AstMountEsp_UnmountDriveW((*iter)->Letter);

			//卸载失败，返回错误代码
			if (!bStatus) return GetLastError();
		}
	}

	return 0;
}

/// <summary>
/// 卸载计算机指定磁盘上全部EFI分区，已经卸载的不会重复卸载
/// </summary>
/// <param name="DiskNumber">将要执行EFI分区卸载操作的磁盘号</param>
/// <param name="LetterList">指向卸载的EFI分区盘符列表的指针，可为nullptr</param>
/// <returns>成功返回0，LetterList存储卸载的EFI分区盘符，磁盘不存在返回ERROR_DRIVE_NOT_EXIST，失败返回GetLastError()返回的错误代码；若返回值为非零值则LetterList为空</returns>
DWORD AstMountEsp_UnMountEfiPartitionW(
	_In_ DWORD DiskNumber,
	_Out_ std::vector<WCHAR>* LetterList
) {
	DWORD dStatus;

	//尝试打开磁盘
	std::wstringstream wss_PhysicalDrivePath;
	wss_PhysicalDrivePath << L"\\\\.\\PhysicalDrive";
	wss_PhysicalDrivePath << DiskNumber;

	HANDLE hDevice = INVALID_HANDLE_VALUE; //硬盘句柄
	//打开磁盘
	hDevice = CreateFile(
		wss_PhysicalDrivePath.str().c_str(),
		GENERIC_READ,
		FILE_SHARE_READ | FILE_SHARE_WRITE,
		NULL,
		OPEN_EXISTING,
		0,
		NULL
	);
	if (hDevice == INVALID_HANDLE_VALUE) return ERROR_DRIVE_NOT_EXIST;

	//获取全部EFI分区信息
	std::vector<PartitionInformationW*> vec_Efi_PIW;
	dStatus = MountEsp_GetEfiPartitionInfoW(&vec_Efi_PIW);
	if (dStatus != 0) { return dStatus; }	//获取失败，返回错误代码

	//遍历列表
	for (auto iter = vec_Efi_PIW.begin(); iter < vec_Efi_PIW.end(); iter++) {
		//查看分区是否已挂载
		if ((*iter)->IsMounted && (*iter)->DiskNumber == DiskNumber) {
			//将卸载分区压入LetterList
			if (LetterList != nullptr) LetterList->push_back((*iter)->Letter);

			//如果挂载，就卸载分区
			BOOL bStatus = AstMountEsp_UnmountDriveW((*iter)->Letter);

			//卸载失败，返回错误代码
			if (!bStatus) return GetLastError();
		}
	}

	return 0;
}

/// <summary>
/// 挂载分区
/// </summary>
/// <param name="Letter">盘符，例如"C","D"</param>
/// <param name="PartionDrivePath">分区路径，如“\\Device\\Harddisk0\\Partition1”</param>
/// <returns>成功返回非0值，失败返回0，可用GetLastError查询错误信息</returns>
BOOL AstMountEsp_MountDriveW(WCHAR Letter, std::wstring PartionDrivePath) {
	//拼接字符串，拼接后字符： X: （X为盘符）
	std::wstring ws;
	ws += Letter;
	ws += L":";

	//挂载分区
	return DefineDosDeviceW(DDD_RAW_TARGET_PATH, ws.c_str(), PartionDrivePath.c_str());
}

/// <summary>
/// 卸载分区
/// </summary>
/// <param name="Letter">盘符，例如"C","D"，字母必须大写</param>
/// <returns>成功返回非0值，失败返回0，可用GetLastError查询错误信息</returns>
BOOL AstMountEsp_UnmountDriveW(WCHAR Letter) {
	//拼接字符串，拼接后字符： X: （X为盘符）
	std::wstring ws;
	ws += Letter;
	ws += L":";

	//卸载分区
	return DefineDosDevice(DDD_REMOVE_DEFINITION, ws.c_str(), nullptr);
}

/// <summary>
///  获取磁盘上所有空闲的盘符
/// </summary>
/// <returns>成功GetLastError()=0，失败GetLastError()为非零值</returns>
std::vector<WCHAR> GetAvaiableLetter() {
	std::vector<WCHAR> DriveLetter;

	//创建完全盘符列表
	std::vector<WCHAR> AllLetter;
	for (WCHAR wc = L'C'; wc <= L'Z'; wc++) {
		std::wstringstream LetterPath;
		LetterPath << L"\\\\.\\" << wc << L":";

		//尝试打开磁盘
		HANDLE hDevice = CreateFile(
			LetterPath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		if (hDevice == INVALID_HANDLE_VALUE) {
			CloseHandle(hDevice);
			DriveLetter.push_back(wc);
		}
		else if (GetLastError() == ERROR_FILE_NOT_FOUND) {
			return (std::vector<WCHAR>)NULL;
		}
	}

	if (DriveLetter.size() == 0) {
		return (std::vector<WCHAR>)NULL;
	}

	SetLastError(0);
	return DriveLetter;
}

/// <summary>
/// 获取计算机全部EFI分区信息
/// </summary>
/// <param name="InfoList">指向存放SimpleEfiPartitionInformationW结构体指针列表的指针</param>
/// <returns>成功返回0，InfoList存储卸载的EFI分区盘符，失败返回GetLastError()返回的错误代码；且InfoList为空</returns>
DWORD AstMontEsp_GetAllEfiPartitionInformationW(
	_Out_ std::vector<SimpleEfiPartitionInformationW*>* InfoList
) {
	DWORD dStatus;

	//获取全部EFI分区信息
	std::vector<PartitionInformationW*> vec_Efi_PIW;
	dStatus = MountEsp_GetEfiPartitionInfoW(&vec_Efi_PIW);
	if (dStatus != 0) { return dStatus; }	//获取失败，返回错误代码

	//开始处理信息
	for (int i = 0; i < vec_Efi_PIW.size(); i++) {
		MallocEx(SimpleEfiPartitionInformationW*, SimpleEfiPIW, 128);

		SimpleEfiPIW->DiskNumber = vec_Efi_PIW[i]->DiskNumber;
		SimpleEfiPIW->PartitionNumber = vec_Efi_PIW[i]->PartitionNumber;
		SimpleEfiPIW->PartitionPath = vec_Efi_PIW[i]->PartitionPath;
		SimpleEfiPIW->PartitionVolumePath = vec_Efi_PIW[i]->PartitionVolumePath;
		SimpleEfiPIW->IsMounted = vec_Efi_PIW[i]->IsMounted;
		SimpleEfiPIW->Letter = vec_Efi_PIW[i]->Letter;

		InfoList->push_back(SimpleEfiPIW);
	}

	return 0;
}

/// <summary>
/// 通过盘符获取指定已挂载EFI分区的信息
/// </summary>
/// <param name="Letter">盘符</param>
/// <param name="PartitionInfo">指向存放分区信息的结构体的指针</param>
/// <returns>获取成功返回0，同时PartitionInfo为指向存放分区信息的结构体的指针；磁盘不存在返回ERROR_DRIVE_NOT_EXIST，失败返回GetLastError()返回的错误代码；获取失败时，PartitionInfo为空</returns>
DWORD AstMontEsp_GetAEfiPartitionInformationW(
	_In_ WCHAR Letter,
	_Out_ SimpleEfiPartitionInformationW* PartitionInfo
) {
	//获取EFI分区信息
	std::vector<SimpleEfiPartitionInformationW*> InfoList;
	DWORD dStatus = AstMontEsp_GetAllEfiPartitionInformationW(&InfoList);
	if (dStatus != 0) return dStatus;

	//开始遍历
	for (auto iter : InfoList) {
		if (iter->Letter == Letter) {
			PartitionInfo = iter;
			return 0;
		}
	}

	//没查找到
	return ERROR_DRIVE_NOT_EXIST;
}
/// <summary>
/// 检查分区是否为ESP分区
/// </summary>
/// <param name="Letter">输入，指向盘符的指针</param>
/// <param name="Status">输出，指向存放检查结果的BOOL变量的指针</param>
/// <returns>成功返回0，同时Status指向存放检查结果的BOOL变量的指针；失败返回GetLastError()返回的错误代码，且不设置Status</returns>
DWORD AstMountEsp_IsPartitionGptW(
	_In_ WCHAR* Letter,
	_Out_ BOOL* Status
) {
	std::vector<PartitionInformationW*> partitionInfoList;		//存放或得到的GPT分区
	DWORD dStatus;												//状态

	//获取计算机上全部GPT分区
	dStatus = MountEsp_GetAllGptPartitionInformationW(
		&partitionInfoList
	);

	if (dStatus != 0) return dStatus;		//获取失败，返回

	for (int i = 0; i < partitionInfoList.size(); i++) {
		auto PIW = partitionInfoList[i];

		if (PIW->IsMounted = true) {
			if ((PIW->Letter == *Letter) && (PIW->PartitionStyle == PARTITION_STYLE_GPT)) {
				*Status = TRUE;
				return 0;
			}
		}
	}

	*Status = FALSE;
	return 0;
}